{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import tensorflow as tf\n",
    "import tensorflow_probability as tfp\n",
    "import scipy.optimize\n",
    "import scipy.io\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import matplotlib.gridspec as gridspec\n",
    "from mpl_toolkits.axes_grid1 import make_axes_locatable\n",
    "import time\n",
    "from pyDOE import lhs         #Latin Hypercube Sampling\n",
    "import datetime, os\n",
    "import seaborn as sns \n",
    "import codecs, json\n",
    "\n",
    "# generates same random numbers each time\n",
    "np.random.seed(1234)\n",
    "tf.random.set_seed(1234)\n",
    "\n",
    "print(\"TensorFlow version: {}\".format(tf.__version__))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "88G3Lt8xn-Oo"
   },
   "source": [
    "# *Data Prep*\n",
    "\n",
    "Training and Testing data is prepared from the solution file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data = scipy.io.loadmat('burgers_shock_mu_01_pi.mat')  \t# Load data from file\n",
    "x = data['x']                                   # 256 points between -1 and 1 [256x1]\n",
    "t = data['t']                                   # 100 time points between 0 and 1 [100x1] \n",
    "usol = data['usol']                             # solution of 256x100 grid points\n",
    "\n",
    "X, T = np.meshgrid(x,t)                         # makes 2 arrays X and T such that u(X[i],T[j])=usol[i][j] are a tuple"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "ZyGxyaOAcqpi"
   },
   "source": [
    "# *Test Data*\n",
    "\n",
    "We prepare the test data to compare against the solution produced by the PINN."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "yddknKA2Xohp"
   },
   "outputs": [],
   "source": [
    "''' X_u_test = [X[i],T[i]] [25600,2] for interpolation'''\n",
    "X_u_test = np.hstack((X.flatten()[:,None], T.flatten()[:,None]))\n",
    "\n",
    "# Domain bounds\n",
    "lb = X_u_test[0]  # [-1. 0.]\n",
    "ub = X_u_test[-1] # [1.  0.99]\n",
    "\n",
    "'''\n",
    "   Fortran Style ('F') flatten,stacked column wise!\n",
    "   u = [c1 \n",
    "        c2\n",
    "        .\n",
    "        .\n",
    "        cn]\n",
    "\n",
    "   u =  [25600x1] \n",
    "'''\n",
    "u = usol.flatten('F')[:,None] "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "aJ5oBRtEXnyu"
   },
   "source": [
    "# *Training Data*\n",
    "\n",
    "The boundary conditions serve as the test data for the PINN and the collocation points are generated using **Latin Hypercube Sampling**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "8UVJmvZbXjXb"
   },
   "outputs": [],
   "source": [
    "def trainingdata(N_u,N_f):\n",
    "\n",
    "    '''Boundary Conditions'''\n",
    "\n",
    "    #Initial Condition -1 =< x =<1 and t = 0  \n",
    "    leftedge_x = np.hstack((X[0,:][:,None], T[0,:][:,None])) #L1\n",
    "    leftedge_u = usol[:,0][:,None]\n",
    "\n",
    "    #Boundary Condition x = -1 and 0 =< t =<1\n",
    "    bottomedge_x = np.hstack((X[:,0][:,None], T[:,0][:,None])) #L2\n",
    "    bottomedge_u = usol[-1,:][:,None]\n",
    "\n",
    "    #Boundary Condition x = 1 and 0 =< t =<1\n",
    "    topedge_x = np.hstack((X[:,-1][:,None], T[:,0][:,None])) #L3\n",
    "    topedge_u = usol[0,:][:,None]\n",
    "\n",
    "    all_X_u_train = np.vstack([leftedge_x, bottomedge_x, topedge_x]) # X_u_train [456,2] (456 = 256(L1)+100(L2)+100(L3))\n",
    "    all_u_train = np.vstack([leftedge_u, bottomedge_u, topedge_u])   #corresponding u [456x1]\n",
    "\n",
    "    #choose random N_u points for training\n",
    "    idx = np.random.choice(all_X_u_train.shape[0], N_u, replace=False) \n",
    "\n",
    "    X_u_train = all_X_u_train[idx, :] #choose indices from  set 'idx' (x,t)\n",
    "    u_train = all_u_train[idx,:]      #choose corresponding u\n",
    "\n",
    "    '''Collocation Points'''\n",
    "\n",
    "    # Latin Hypercube sampling for collocation points \n",
    "    # N_f sets of tuples(x,t)\n",
    "    X_f_train = lb + (ub-lb)*lhs(2,N_f) \n",
    "    X_f_train = np.vstack((X_f_train, X_u_train)) # append training points to collocation points \n",
    "\n",
    "    return X_f_train, X_u_train, u_train \n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "dp4nc2S7bwzz"
   },
   "source": [
    "# **PINN**\n",
    "\n",
    "Generate a **PINN** of L hidden layers, each with n neurons. \n",
    "\n",
    "Initialization: ***Xavier***\n",
    "\n",
    "Activation: *tanh (x)*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "class Sequentialmodel(tf.Module): \n",
    "    def __init__(self, layers, name=None):\n",
    "       \n",
    "        self.W = []  #Weights and biases\n",
    "        self.parameters = 0 #total number of parameters\n",
    "        \n",
    "        for i in range(len(layers)-1):\n",
    "            \n",
    "            input_dim = layers[i]\n",
    "            output_dim = layers[i+1]\n",
    "            \n",
    "            #Xavier standard deviation \n",
    "            std_dv = np.sqrt((2.0/(input_dim + output_dim)))\n",
    "\n",
    "            #weights = normal distribution * Xavier standard deviation + 0\n",
    "            w = tf.random.normal([input_dim, output_dim], dtype = 'float64') * std_dv\n",
    "                       \n",
    "            w = tf.Variable(w, trainable=True, name = 'w' + str(i+1))\n",
    "\n",
    "            b = tf.Variable(tf.cast(tf.zeros([output_dim]), dtype = 'float64'), trainable = True, name = 'b' + str(i+1))\n",
    "                    \n",
    "            self.W.append(w)\n",
    "            self.W.append(b)\n",
    "            \n",
    "            self.parameters +=  input_dim * output_dim + output_dim\n",
    "    \n",
    "    def evaluate(self,x):\n",
    "        \n",
    "        a = x\n",
    "        \n",
    "        for i in range(len(layers)-2):\n",
    "            \n",
    "            z = tf.add(tf.matmul(a, self.W[2*i]), self.W[2*i+1])\n",
    "            a = tf.nn.tanh(z)\n",
    "            \n",
    "        a = tf.add(tf.matmul(a, self.W[-2]), self.W[-1]) # For regression, no activation to last layer\n",
    "        return a\n",
    "    \n",
    "    def get_weights(self):\n",
    "\n",
    "        parameters_1d = []  # [.... W_i,b_i.....  ] 1d array\n",
    "        \n",
    "        for i in range (len(layers)-1):\n",
    "            \n",
    "            w_1d = tf.reshape(self.W[2*i],[-1])   #flatten weights \n",
    "            b_1d = tf.reshape(self.W[2*i+1],[-1]) #flatten biases\n",
    "            \n",
    "            parameters_1d = tf.concat([parameters_1d, w_1d], 0) #concat weights \n",
    "            parameters_1d = tf.concat([parameters_1d, b_1d], 0) #concat biases\n",
    "        \n",
    "        return parameters_1d\n",
    "        \n",
    "    def set_weights(self,parameters):\n",
    "                \n",
    "        for i in range (len(layers)-1):\n",
    "\n",
    "            shape_w = tf.shape(self.W[2*i]).numpy() # shape of the weight tensor\n",
    "            size_w = tf.size(self.W[2*i]).numpy() #size of the weight tensor \n",
    "            \n",
    "            shape_b = tf.shape(self.W[2*i+1]).numpy() # shape of the bias tensor\n",
    "            size_b = tf.size(self.W[2*i+1]).numpy() #size of the bias tensor \n",
    "                        \n",
    "            pick_w = parameters[0:size_w] #pick the weights \n",
    "            self.W[2*i].assign(tf.reshape(pick_w,shape_w)) # assign  \n",
    "            parameters = np.delete(parameters,np.arange(size_w),0) #delete \n",
    "            \n",
    "            pick_b = parameters[0:size_b] #pick the biases \n",
    "            self.W[2*i+1].assign(tf.reshape(pick_b,shape_b)) # assign \n",
    "            parameters = np.delete(parameters,np.arange(size_b),0) #delete \n",
    "\n",
    "            \n",
    "    def loss_BC(self,x,y):\n",
    "\n",
    "        loss_u = tf.reduce_mean(tf.square(y-self.evaluate(x)))\n",
    "        return loss_u\n",
    "\n",
    "    def loss_PDE(self, x_to_train_f):\n",
    "    \n",
    "        g = tf.Variable(x_to_train_f, dtype = 'float64')\n",
    "    \n",
    "        nu = 0.01/np.pi\n",
    "\n",
    "        x_f = g[:,0:1]\n",
    "        t_f = g[:,1:2]\n",
    "\n",
    "        with tf.GradientTape(persistent=True) as tape:\n",
    "\n",
    "            tape.watch(x_f)\n",
    "            tape.watch(t_f)\n",
    "\n",
    "            g = tf.stack([x_f[:,0], t_f[:,0]], axis=1)   \n",
    "\n",
    "            z = self.evaluate(g)\n",
    "            u_x = tape.gradient(z,x_f)\n",
    "\n",
    "        u_t = tape.gradient(z,t_f)    \n",
    "        u_xx = tape.gradient(u_x, x_f)\n",
    "\n",
    "        del tape\n",
    "\n",
    "        f = u_t + (self.evaluate(g))*(u_x) - (nu)*u_xx\n",
    "\n",
    "        loss_f = tf.reduce_mean(tf.square(f))\n",
    "\n",
    "        return loss_f\n",
    "    \n",
    "    def loss(self,x,y,g):\n",
    "\n",
    "        loss_u = self.loss_BC(x,y)\n",
    "        loss_f = self.loss_PDE(g)\n",
    "\n",
    "        loss = loss_u + loss_f\n",
    "\n",
    "        return loss\n",
    "    \n",
    "    def optimizerfunc(self,parameters):\n",
    "        \n",
    "        self.set_weights(parameters)\n",
    "       \n",
    "        with tf.GradientTape() as tape:\n",
    "            tape.watch(self.trainable_variables)\n",
    "            \n",
    "            loss_val = self.loss(X_u_train, u_train, X_f_train)\n",
    "            \n",
    "        grads = tape.gradient(loss_val,self.trainable_variables)\n",
    "                \n",
    "        del tape\n",
    "        \n",
    "        grads_1d = [ ] #flatten grads \n",
    "        \n",
    "        for i in range (len(layers)-1):\n",
    "\n",
    "            grads_w_1d = tf.reshape(grads[2*i],[-1]) #flatten weights \n",
    "            grads_b_1d = tf.reshape(grads[2*i+1],[-1]) #flatten biases\n",
    "\n",
    "            grads_1d = tf.concat([grads_1d, grads_w_1d], 0) #concat grad_weights \n",
    "            grads_1d = tf.concat([grads_1d, grads_b_1d], 0) #concat grad_biases\n",
    "\n",
    "        return loss_val, grads_1d\n",
    "    \n",
    "    def optimizer_callback(self,parameters):\n",
    "        \n",
    "        self.set_weights(parameters)\n",
    "        \n",
    "        loss_value = self.loss(X_u_train, u_train, X_f_train)\n",
    "        \n",
    "        u_pred = self.evaluate(X_u_test)\n",
    "        error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)\n",
    "        \n",
    "        tf.print(\"loss:\", loss_value, \"error:\", error_vec)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "bOjuHdzAhib-"
   },
   "source": [
    "# *Solution Plot*"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "colab": {},
    "colab_type": "code",
    "id": "UWqNuRMLhg4m"
   },
   "outputs": [],
   "source": [
    "def solutionplot(u_pred,X_u_train,u_train):\n",
    "    \n",
    "    fig, ax = plt.subplots()\n",
    "    ax.axis('off')\n",
    "\n",
    "    gs0 = gridspec.GridSpec(1, 2)\n",
    "    gs0.update(top=1-0.06, bottom=1-1/3, left=0.15, right=0.85, wspace=0)\n",
    "    ax = plt.subplot(gs0[:, :])\n",
    "\n",
    "    h = ax.imshow(u_pred, interpolation='nearest', cmap='rainbow', \n",
    "                extent=[T.min(), T.max(), X.min(), X.max()], \n",
    "                origin='lower', aspect='auto')\n",
    "    divider = make_axes_locatable(ax)\n",
    "    cax = divider.append_axes(\"right\", size=\"5%\", pad=0.05)\n",
    "    fig.colorbar(h, cax=cax)\n",
    "    \n",
    "    ax.plot(X_u_train[:,1], X_u_train[:,0], 'kx', label = 'Data (%d points)' % (u_train.shape[0]), markersize = 4, clip_on = False)\n",
    "\n",
    "    line = np.linspace(x.min(), x.max(), 2)[:,None]\n",
    "    ax.plot(t[25]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "    ax.plot(t[50]*np.ones((2,1)), line, 'w-', linewidth = 1)\n",
    "    ax.plot(t[75]*np.ones((2,1)), line, 'w-', linewidth = 1)    \n",
    "\n",
    "    ax.set_xlabel('$t$')\n",
    "    ax.set_ylabel('$x$')\n",
    "    ax.legend(frameon=False, loc = 'best')\n",
    "    ax.set_title('$u(x,t)$', fontsize = 10)\n",
    "    \n",
    "    ''' \n",
    "    Slices of the solution at points t = 0.25, t = 0.50 and t = 0.75\n",
    "    '''\n",
    "    \n",
    "    ####### Row 1: u(t,x) slices ##################\n",
    "    gs1 = gridspec.GridSpec(1, 3)\n",
    "    gs1.update(top=1-1/3, bottom=0, left=0.1, right=0.9, wspace=0.5)\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 0])\n",
    "    ax.plot(x,usol.T[25,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[25,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')    \n",
    "    ax.set_title('$t = 0.25s$', fontsize = 10)\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 1])\n",
    "    ax.plot(x,usol.T[50,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[50,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])\n",
    "    ax.set_title('$t = 0.50s$', fontsize = 10)\n",
    "    ax.legend(loc='upper center', bbox_to_anchor=(0.5, -0.35), ncol=5, frameon=False)\n",
    "\n",
    "    ax = plt.subplot(gs1[0, 2])\n",
    "    ax.plot(x,usol.T[75,:], 'b-', linewidth = 2, label = 'Exact')       \n",
    "    ax.plot(x,u_pred.T[75,:], 'r--', linewidth = 2, label = 'Prediction')\n",
    "    ax.set_xlabel('$x$')\n",
    "    ax.set_ylabel('$u(x,t)$')\n",
    "    ax.axis('square')\n",
    "    ax.set_xlim([-1.1,1.1])\n",
    "    ax.set_ylim([-1.1,1.1])    \n",
    "    ax.set_title('$t = 0.75s$', fontsize = 10)\n",
    "    \n",
    "    plt.savefig('Burgers.png',dpi = 500)   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "colab_type": "text",
    "id": "YRuuEXx-eeWa"
   },
   "source": [
    "# *Model Training and Testing*\n",
    "\n",
    "A function '**model**' is defined to generate a NN as per the input set of hyperparameters, which is then trained and tested. The L2 Norm of the solution error is returned as a comparison metric"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "ename": "ImportError",
     "evalue": "This version of TensorFlow Probability requires TensorFlow version >= 2.3; Detected an installation of version 2.0.0. Please upgrade TensorFlow to proceed.",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mImportError\u001b[0m                               Traceback (most recent call last)",
      "\u001b[1;32m<ipython-input-2-41494c8c96ff>\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[1;32m----> 1\u001b[1;33m \u001b[1;32mimport\u001b[0m \u001b[0mtensorflow_probability\u001b[0m \u001b[1;32mas\u001b[0m \u001b[0mtfp\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m",
      "\u001b[1;32mc:\\anaconda\\envs\\tf2\\lib\\site-packages\\tensorflow_probability\\__init__.py\u001b[0m in \u001b[0;36m<module>\u001b[1;34m\u001b[0m\n\u001b[0;32m     67\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     68\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 69\u001b[1;33m \u001b[0m_ensure_tf_install\u001b[0m\u001b[1;33m(\u001b[0m\u001b[1;33m)\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0m\u001b[0;32m     70\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     71\u001b[0m \u001b[1;31m# Cleanup symbols to avoid polluting namespace.\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;32mc:\\anaconda\\envs\\tf2\\lib\\site-packages\\tensorflow_probability\\__init__.py\u001b[0m in \u001b[0;36m_ensure_tf_install\u001b[1;34m()\u001b[0m\n\u001b[0;32m     64\u001b[0m         \"Please upgrade TensorFlow to proceed.\".format(\n\u001b[0;32m     65\u001b[0m             \u001b[0mrequired\u001b[0m\u001b[1;33m=\u001b[0m\u001b[0mrequired_tensorflow_version\u001b[0m\u001b[1;33m,\u001b[0m\u001b[1;33m\u001b[0m\u001b[1;33m\u001b[0m\u001b[0m\n\u001b[1;32m---> 66\u001b[1;33m             present=tf.__version__))\n\u001b[0m\u001b[0;32m     67\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n\u001b[0;32m     68\u001b[0m \u001b[1;33m\u001b[0m\u001b[0m\n",
      "\u001b[1;31mImportError\u001b[0m: This version of TensorFlow Probability requires TensorFlow version >= 2.3; Detected an installation of version 2.0.0. Please upgrade TensorFlow to proceed."
     ]
    }
   ],
   "source": [
    "N_u = 100 #Total number of data points for 'u'\n",
    "N_f = 10000 #Total number of collocation points \n",
    "\n",
    "# Training data\n",
    "X_f_train, X_u_train, u_train = trainingdata(N_u,N_f)\n",
    "\n",
    "layers = np.array([2,20,20,20,20,20,20,20,20,1]) #8 hidden layers\n",
    "\n",
    "PINN = Sequentialmodel(layers)\n",
    "\n",
    "init_params = PINN.get_weights().numpy()\n",
    "\n",
    "start_time = time.time() \n",
    "\n",
    "# train the model with TFP L-BFGS optimizer\n",
    "results = tfp.optimizer.lbfgs_minimize(PINN.optimizerfunc,\n",
    "                                       init_params,\n",
    "                                       num_correction_pairs=200,\n",
    "                                       tolerance = 1*np.finfo(float).eps, \n",
    "                                       x_tolerance = 0,\n",
    "                                       f_relative_tolerance = 0,\n",
    "                                       max_iterations = 5000,\n",
    "                                       parallel_iterations = 1,\n",
    "                                       max_line_search_iterations = 50)\n",
    "\n",
    "elapsed = time.time() - start_time                \n",
    "print('Training time: %.2f' % (elapsed))\n",
    "\n",
    "print(results)\n",
    "\n",
    "PINN.set_weights(results.position)\n",
    "\n",
    "''' Model Accuracy ''' \n",
    "u_pred = PINN.evaluate(X_u_test)\n",
    "\n",
    "error_vec = np.linalg.norm((u-u_pred),2)/np.linalg.norm(u,2)        # Relative L2 Norm of the error (Vector)\n",
    "print('Test Error: %.5f'  % (error_vec))\n",
    "\n",
    "u_pred = np.reshape(u_pred,(256,100),order='F')                        # Fortran Style ,stacked column wise!\n",
    "\n",
    "''' Solution Plot '''\n",
    "solutionplot(u_pred,X_u_train,u_train)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# data = np.loadtxt(\"original_run.txt\", comments=\"#\", delimiter=\" \", unpack=False)\n",
    "\n",
    "# # sns.set()\n",
    "\n",
    "# fig, ax = plt.subplots()\n",
    "# plt.semilogy(np.arange(1,5001),data[:,0], label = 'Loss')\n",
    "# plt.semilogy(np.arange(1,5001),data[:,1], label = 'Error')\n",
    "\n",
    "# plt.legend()\n",
    "# ax.set_xlabel('iterations')\n",
    "# ax.set_ylabel('Loss/Error')\n",
    "# fig.suptitle('Convergence Plot')\n",
    "# plt.grid()\n",
    "# plt.show()\n",
    "\n",
    "# # fig.savefig('original_run_history.png', dpi = 500)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
